##
# 用于将整幅图片中的所有标注的汽车抠出来并归一化尺寸存放
##
import os

from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from data.datasets import CarDataset, ToTensor
import cv2
import numpy as np

BATCH_SIZE = 1
res_path = 'E://workspace//Datasets//Car//img//'
anno_path = 'E://workspace//Datasets//Car//anno.txt'

# 加载数据
train_datasets = CarDataset(anno_file='E://workspace//UA-DETRAC//train_gt.txt',
                           root_dir='E://workspace//UA-DETRAC//DETRAC-train-data//Insight-MVT_Annotation_Train//')
train_loader = DataLoader(train_datasets, batch_size=BATCH_SIZE,
                        shuffle=True)# 加入num_workers=4会报错，这是Pytorch在windows下的BUG

idx = 0
# 用于存放车辆图像的标签的
f = open(anno_path, 'w')
for i_batch, sample_batched in enumerate(train_loader):
    image = sample_batched['image'].cpu().data.numpy()[0]
    bboxes = sample_batched['bbox_target'].cpu().data.numpy().squeeze()

    if len(bboxes[:]) == 4 and isinstance(bboxes[0], np.int32):# 考虑只有一个标注框的情形
        [x1, y1, x2, y2] = bboxes
        crop_img = image[y1:y2, x1:x2, :]
        crop_img = cv2.resize(crop_img, (32, 32))

        save_file = os.path.join(res_path, "%s.jpg" % idx)
        f.write(save_file + ' 1\n')
        print(cv2.imwrite(save_file, crop_img))
        idx += 1

        # cv2.imshow("", crop_img)
        # cv2.waitKey(100)
    else:
        for bbox in bboxes:
            [x1, y1, x2, y2] = bbox
            crop_img = image[y1:y2, x1:x2, :]
            crop_img = cv2.resize(crop_img, (32, 32))

            save_file = os.path.join(res_path, "%s.jpg" % idx)
            f.write(save_file + ' 1\n')
            print(cv2.imwrite(save_file, crop_img))
            idx += 1

            # cv2.imshow("", crop_img)
            # cv2.waitKey(100)

    # print(i_batch, sample_batched['image'].size(),
    #       sample_batched['bbox_target'])